library(yaml)
library(data.table)
library(tidyverse)
library(foreign)
library(plyr)

calc_partisan_affect_polarization  <- function(db, non_na_affect){

  # Create column for each party's size
  sizes_party            <- aggregate(db$weight, by = list(db$party), FUN = sum)
  colnames(sizes_party)  <- c("party", "party_size")
  for (i in non_na_affect){
    db[, sprintf("party_size%s", i)] <- sizes_party[which(sizes_party$party == i), "party_size"]
  }
  
  # Create column for each individual's party's weight
  b                   <- cbind(seq_along(db$party), match(sprintf("party_size%s", db$party), colnames(db)))
  db$own_party_weight <- db[b]
  db$own_party_weight <- as.numeric(db$own_party_weight)
  
  # Construct relative weights columns taking into account NAs
  db$total_rel_weight <- rowSums(db[, sprintf("party_size%s", non_na_affect)] * 
                                   !is.na(db[, sprintf("feel_party%s", non_na_affect)]))
  for(i in non_na_affect){
    db[, sprintf("rel_party_weight%s", i)] <- db[, sprintf("party_size%s", i)] / (db$total_rel_weight - db$own_party_weight)
  }
  
  db$own_party_rel_weight <- db$own_party_weight / (db$total_rel_weight - db$own_party_weight)
  
  # Weighted sum over parties for each individual
  db$pi_in  <- db$own_party_feel
  db$pi_out <- row_Sums(db[, sprintf("feel_party%s", non_na_affect)] * db[, sprintf("rel_party_weight%s", non_na_affect)]) - 
                              (db$own_party_feel * db$own_party_rel_weight) # Need to subtract own-party contribution to out-party feelings since taking a row_sum
  db$pi     <- db$pi_in - db$pi_out
  
  # Weighted sum across individuals 
  if (sum(is.na(db$pi)) > 0) stop("NA in pi is an error.")
  
  affect     <- sum(db$pi     * db$weight) / sum(db$weight)
  affect_in  <- sum(db$pi_in  * db$weight) / sum(db$weight)
  affect_out <- sum(db$pi_out * db$weight) / sum(db$weight)
  
  sample_sd  <- sqrt(sum(db$weight * (db$pi - affect)^2)  / sum(db$weight)) # https://en.wikipedia.org/wiki/Weighted_arithmetic_mean#Variance_weights
  sd_aff     <- sqrt(sum( (db$weight / sum(db$weight))^2 * (db$pi - affect)^2 ))
  return(list("affect" = affect, "affect_sd" = sd_aff, "sample_sd" = sample_sd, 
              "affect_in" = affect_in, "affect_out" = affect_out))
}


row_Sums <- function(x){
  temp         <- rowSums(x, na.rm = T)
  all_na       <- rowSums(is.na(x)) == ncol(x)
  temp[all_na] <- NA
  return(temp)
}


write_gslab_table <- function(TABLE, output, header){
  write.table(TABLE, file = output, row.names = FALSE, col.names = FALSE, sep = "\t")
  table <- read_file(output) # http://stackoverflow.com/questions/9068397/import-text-file-as-single-character-string
  table <- sprintf("%s\n%s", header, gsub('"', '', table))
  write_file(table, output)
}


load_worldbank_data <- function(varname, path, countries, source=FALSE) {
  
  data <- read.csv(path, skip=4) %>%
          reshape(varying = sprintf("X%s", 1960:2020), v.names = c("X"), times = c(1960:2020), direction = "long") %>%
          select("Country.Name", "X", "time") %>%
          filter(Country.Name %in% countries) %>%
          drop_na()
  
  if (source == TRUE) {
    data        <- data %>% mutate(source = "World Bank")
    names(data) <- c("country", varname, "years", "source")
  } else {
    names(data) <- c("country", varname, "years")
  }
  
  return(data)
}


table_reg_driver_year <- function(driver_data, countries, driver){
  
  TAB  <- c()
  for (c in countries){
    temp <- subset(driver_data, country == c & years >= 1970)
    
    if (nrow(temp) > 1) {
      out  <- summary(lm(temp[, driver] ~ temp[, "years"]))
      TAB  <- rbind(TAB, as.matrix(out$coefficients[2, c(1:2)]))
    } else {
      TAB  <- rbind(TAB, as.matrix(c(NA,NA)))
    }
  }
  
  return(TAB)
}

